use anyhow::Result;
use std::net::SocketAddr;
use tokio::net::UdpSocket as TokioUdpSocket;
use tracing::{debug, error, info};
use std::time::Duration;

#[derive(Debug)]
pub struct RtpHandler {
    socket: TokioUdpSocket,
    local_addr: SocketAddr,
    remote_addr: Option<SocketAddr>,
    sequence_number: u16,
    timestamp: u32,
    ssrc: u32,
}

#[derive(Debug)]
pub struct RtpPacket {
    pub version: u8,
    pub padding: bool,
    pub extension: bool,
    pub csrc_count: u8,
    pub marker: bool,
    pub payload_type: u8,
    pub sequence_number: u16,
    pub timestamp: u32,
    pub ssrc: u32,
    pub payload: Vec<u8>,
}

impl RtpHandler {
    pub async fn new(local_port: u16) -> Result<Self> {
        let addr = format!("0.0.0.0:{}", local_port);
        let socket = TokioUdpSocket::bind(&addr).await?;
        let local_addr = socket.local_addr()?;
        
        info!("RTP处理器启动在端口: {}", local_port);
        
        Ok(Self {
            socket,
            local_addr,
            remote_addr: None,
            sequence_number: 0,
            timestamp: 0,
            ssrc: rand::random(),
        })
    }

    pub fn set_remote_addr(&mut self, addr: SocketAddr) {
        self.remote_addr = Some(addr);
        info!("设置RTP远程地址: {}", addr);
    }

    pub async fn send_audio_packet(&mut self, audio_data: &[u8]) -> Result<()> {
        if let Some(remote_addr) = self.remote_addr {
            let packet = self.create_rtp_packet(audio_data);
            let packet_data = self.serialize_rtp_packet(&packet)?;
            
            self.socket.send_to(&packet_data, remote_addr).await?;
            
            // 更新序列号和时间戳
            self.sequence_number = self.sequence_number.wrapping_add(1);
            self.timestamp = self.timestamp.wrapping_add(160); // 20ms at 8kHz
            
            debug!("发送RTP包: seq={}, ts={}", packet.sequence_number, packet.timestamp);
        }
        
        Ok(())
    }

    pub async fn receive_audio_packet(&self) -> Result<Option<RtpPacket>> {
        let mut buffer = [0; 1500];
        
        // 使用超时来避免阻塞
        match tokio::time::timeout(Duration::from_millis(1), self.socket.recv_from(&mut buffer)).await {
            Ok(Ok((n, addr))) => {
                if let Some(remote_addr) = self.remote_addr {
                    if addr == remote_addr {
                        let packet = self.deserialize_rtp_packet(&buffer[..n])?;
                        debug!("接收RTP包: seq={}, ts={}", packet.sequence_number, packet.timestamp);
                        return Ok(Some(packet));
                    }
                }
            }
            Ok(Err(e)) => {
                error!("接收RTP包时出错: {}", e);
            }
            Err(_) => {
                // 超时，没有数据可读
            }
        }
        
        Ok(None)
    }

    fn create_rtp_packet(&self, payload: &[u8]) -> RtpPacket {
        RtpPacket {
            version: 2,
            padding: false,
            extension: false,
            csrc_count: 0,
            marker: false,
            payload_type: 0, // PCMU
            sequence_number: self.sequence_number,
            timestamp: self.timestamp,
            ssrc: self.ssrc,
            payload: payload.to_vec(),
        }
    }

    fn serialize_rtp_packet(&self, packet: &RtpPacket) -> Result<Vec<u8>> {
        let mut data = Vec::new();
        
        // RTP头部 (12字节)
        let first_byte = (packet.version << 6) | 
                        ((packet.padding as u8) << 5) | 
                        ((packet.extension as u8) << 4) | 
                        packet.csrc_count;
        data.push(first_byte);
        
        let second_byte = ((packet.marker as u8) << 7) | packet.payload_type;
        data.push(second_byte);
        
        // 序列号 (2字节)
        data.extend_from_slice(&packet.sequence_number.to_be_bytes());
        
        // 时间戳 (4字节)
        data.extend_from_slice(&packet.timestamp.to_be_bytes());
        
        // SSRC (4字节)
        data.extend_from_slice(&packet.ssrc.to_be_bytes());
        
        // 负载
        data.extend_from_slice(&packet.payload);
        
        Ok(data)
    }

    fn deserialize_rtp_packet(&self, data: &[u8]) -> Result<RtpPacket> {
        if data.len() < 12 {
            return Err(anyhow::anyhow!("RTP包太小"));
        }
        
        let first_byte = data[0];
        let second_byte = data[1];
        
        let version = first_byte >> 6;
        let padding = (first_byte >> 5) & 1 != 0;
        let extension = (first_byte >> 4) & 1 != 0;
        let csrc_count = first_byte & 0x0f;
        
        let marker = (second_byte >> 7) != 0;
        let payload_type = second_byte & 0x7f;
        
        let sequence_number = u16::from_be_bytes([data[2], data[3]]);
        let timestamp = u32::from_be_bytes([data[4], data[5], data[6], data[7]]);
        let ssrc = u32::from_be_bytes([data[8], data[9], data[10], data[11]]);
        
        let payload = data[12..].to_vec();
        
        Ok(RtpPacket {
            version,
            padding,
            extension,
            csrc_count,
            marker,
            payload_type,
            sequence_number,
            timestamp,
            ssrc,
            payload,
        })
    }
}

#[derive(Debug)]
pub struct AudioProcessor {
    rtp_handler: RtpHandler,
    sample_rate: u32,
    channels: u16,
}

impl AudioProcessor {
    pub async fn new(rtp_port: u16, sample_rate: u32, channels: u16) -> Result<Self> {
        let rtp_handler = RtpHandler::new(rtp_port).await?;
        
        Ok(Self {
            rtp_handler,
            sample_rate,
            channels,
        })
    }

    pub async fn start_audio_session(&mut self, remote_addr: SocketAddr) -> Result<()> {
        self.rtp_handler.set_remote_addr(remote_addr);
        info!("音频会话已启动");
        Ok(())
    }

    pub async fn send_audio(&mut self, audio_data: &[u8]) -> Result<()> {
        self.rtp_handler.send_audio_packet(audio_data).await
    }

    pub async fn receive_audio(&self) -> Result<Option<Vec<u8>>> {
        if let Some(packet) = self.rtp_handler.receive_audio_packet().await? {
            return Ok(Some(packet.payload));
        }
        Ok(None)
    }
} 